import os
import shutil
import random
from tqdm import tqdm

# 原始数据路径（手动放入）
SOURCE_DIR = "data/train"
OUTPUT_DIR = "data/processed"

# 划分比例
train_ratio = 0.8

os.makedirs(f"{OUTPUT_DIR}/train/cats", exist_ok=True)
os.makedirs(f"{OUTPUT_DIR}/train/dogs", exist_ok=True)
os.makedirs(f"{OUTPUT_DIR}/val/cats", exist_ok=True)
os.makedirs(f"{OUTPUT_DIR}/val/dogs", exist_ok=True)

cats = [f for f in os.listdir(SOURCE_DIR) if f.startswith("cat")]
dogs = [f for f in os.listdir(SOURCE_DIR) if f.startswith("dog")]

for label, files in [("cats", cats), ("dogs", dogs)]:
    random.shuffle(files)
    split = int(len(files) * train_ratio)
    train_files, val_files = files[:split], files[split:]

    for f in tqdm(train_files, desc=f"Copying train {label}"):
        shutil.copy(os.path.join(SOURCE_DIR, f), f"{OUTPUT_DIR}/train/{label}/{f}")

    for f in tqdm(val_files, desc=f"Copying val {label}"):
        shutil.copy(os.path.join(SOURCE_DIR, f), f"{OUTPUT_DIR}/val/{label}/{f}")

print("✅ 数据集划分完成！")
